from time import sleep
from mss.windows import MSS as mss
import multiprocessing
import pyautogui
import cv2
import numpy as np
from os import remove
from PIL import ImageGrab


class CaptureScreen:

    def __init__(self, video_path: str) -> None:
        self.fps = 20
        self.video_path = video_path.replace('\\', '/') + '.avi'

        self.process = multiprocessing.Process(
            target=self.capture_screen,
            args=(),
            name='Screen Capture'
        )

    def capture_screen(self):

        SCREEN_SIZE = tuple(pyautogui.size())
        fourcc = cv2.VideoWriter_fourcc(*'XVID')
        out = cv2.VideoWriter(self.video_path, fourcc, self.fps, (SCREEN_SIZE))
        img = None
        print(f'Screen capture started')

        # MSS
        w, h = pyautogui.size()
        # (left, top, width, height)
        region = (0, 0, w, h)
        with mss() as sct:
            while True:
                try:
                    try:
                        # = sct.grab(( 0,0,w,h) ) 
                        img = pyautogui.screenshot()
                        # img = ImageGrab.grab()
                        frame = np.array(img)
                        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                        out.write(frame)
                    except IOError:
                        pass
                except KeyboardInterrupt:
                    break
            cv2.destroyAllWindows()
            out.release()

    def start_capture(self):
        self.process.start()

    def finish_capture(self):
        if not self.process.is_alive():
            print('The process is not executing')
        else:
            self.process.terminate()
            sleep(3)
            print('Screen capture finished')


if __name__ == '__main__':
    c = CaptureScreen(video_path=r'E:\\game\\python-learn\\temp\\output')
    c.start_capture()
    sleep(15)
    c.finish_capture()
